package com.br.aiops.receiver.es.query;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Maps;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

@Component
public class Query {

    private Logger LOGGER = LoggerFactory.getLogger(Query.class);

    public static final String FIELD_HITS = "hits";
    public static final String FIELD_TOTAL = "total";
    public static final String FIELD_VALUE = "value";
    public static final String FIELD_SOURCE = "_source";
    public static final String FIELD_BEAT = "beat";
    public static final String FIELD_NAME = "name";
    public static final String FIELD_MESSAGE = "message";

    @Autowired
    private RestClient client;

    public Result query(DSL dsl, long start, long end) throws IOException {
        Request request = new Request(dsl.getMethod(), dsl.getEndpoint());
        request.setJsonEntity(dsl.getReplacedBody(start, end));
        LOGGER.info("发送请求数据：{}.", dsl);
        Response response = client.performRequest(request);
        String responseBody = EntityUtils.toString(response.getEntity());
        JSONObject object = JSONObject.parseObject(responseBody);
        LOGGER.info("返回结果：{}.", object);
        JSONArray hits = object.getJSONObject(FIELD_HITS).getJSONArray(FIELD_HITS);
        List<Record> records = hits.stream().map(obj -> buildRecord((JSONObject) obj, dsl.getResultFields()))
                .collect(Collectors.toList());
        Set<String> hosts = getHosts(hits);
        Long count = getCount(object);
        return Result.builder().records(records).totalCount(count).hosts(hosts).message(getMessage(hits)).build();
    }

    private Long getCount(JSONObject object) {
        Object total = object.getJSONObject(FIELD_HITS).get(FIELD_TOTAL);
        if (total instanceof JSONObject) {
            return object.getJSONObject(FIELD_HITS).getJSONObject(FIELD_TOTAL).getLong(FIELD_VALUE);
        }
        return object.getJSONObject(FIELD_HITS).getLong(FIELD_TOTAL);
    }

    private Set<String> getHosts(JSONArray hits) {
        return hits.stream().map(obj -> getHost(((JSONObject)obj).getJSONObject(FIELD_SOURCE))).filter(Objects::nonNull).collect(Collectors.toSet());
    }

    private String getMessage(JSONArray hits) {
        return hits.stream().findFirst().map(obj -> ((JSONObject)obj).getJSONObject(FIELD_SOURCE).getString(FIELD_MESSAGE)).orElse("es没有message字段。");
    }

    private String getHost(JSONObject source) {
        if (source.containsKey(FIELD_BEAT) && source.getJSONObject(FIELD_BEAT).containsKey(FIELD_NAME)) {
            return source.getJSONObject(FIELD_BEAT).getString(FIELD_NAME);
        }
        return null;
    }

    private Record buildRecord(JSONObject obj, List<String> resultFields) {
        Map<String, Object> fields = Maps.newHashMap();
        resultFields.stream().filter(field -> obj.getJSONObject(FIELD_SOURCE).containsKey(field)).forEach(field -> fields.put(field, obj.getJSONObject(FIELD_SOURCE).getString(field)));
        return Record.builder().fields(fields).build();
    }
}
